import argparse
import datetime
import os

from lib.config import cfg, update_config
from lib.core import Trainer
from lib.core import Tester
import numpy as np
import torch
import random
from lib.dataset import DATASET


def same_seeds(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


if __name__ == '__main__':
    os.environ['TORCH_DISTRIBUTED_DEBUG']='DETAIL'

    same_seeds(90)

    parser = argparse.ArgumentParser()

    parser.add_argument('--cf',
                        default='config/baseline_cub200_conv4_5way5shot15query.yaml',
                        dest='config_file',
                        help="the customized config file, default is 'default.yaml' if None",
                        required=False)

    parser.add_argument('--opts',
                        help="modify config options using the command-line",
                        nargs='*',
                        required=False)

    args = parser.parse_args()
    update_config(cfg, args)

    tester = Tester(cfg)
    tester.run()
